import time
from playwright.sync_api import sync_playwright

def wait_for_all_images(page):
    page.evaluate("""
        async () => {
            const images = Array.from(document.images);
            await Promise.all(images.map(img => {
                if (img.complete) return Promise.resolve();
                return new Promise((resolve, reject) => {
                    img.onload = resolve;
                    img.onerror = reject;
                });
            }));
        }
    """)

def scroll_to_bottom(page):
    page.evaluate("""
        async () => {
            await new Promise(resolve => {
                let totalHeight = 0;
                const distance = 100;
                const timer = setInterval(() => {
                    window.scrollBy(0, distance);
                    totalHeight += distance;
                    if (totalHeight >= document.body.scrollHeight) {
                        clearInterval(timer);
                        resolve();
                    }
                }, 100);
            });
        }
    """)

if __name__ == "__main__":
    url = input("请输入网页地址: ")
    output_path = "screenshot.png"
    with sync_playwright() as p:
        browser = p.chromium.launch(headless=False)
        page = browser.new_page()
        try:
            page.goto(url, wait_until="load", timeout=30000)  # 增大超时时间，等待页面加载即可
        except Exception as e:
            print(f"页面加载失败: {e}")
            browser.close()
            exit(1)
        scroll_to_bottom(page)
        wait_for_all_images(page)
        time.sleep(1)
        page.screenshot(path=output_path, full_page=True)
        browser.close()
    print(f"网页截图已保存为 {output_path}")
